import shutil
import os
import pprint
import json
import re
import time
try:
    from gradio_client import Client, handle_file
except ImportError:
    # 如果无法导入gradio_client，提示用户安装
    raise ImportError("请先安装gradio_client模块。可以使用命令: pip install gradio_client")

class MyAudioClient:
    def __init__(self, config_path='config.json'):
        self.config = self.load_config(config_path)
        self.workspace_path = os.path.abspath(os.path.join(os.getcwd(), ".."))
        self.folder_name = self.config['folder_name']
        self.folder_path = f'{self.workspace_path}/{self.folder_name}'
        self.audio_path = f'{self.folder_path}/audio'
        self.txt_name = self.config['txt_name']
        self.file_full_path = f'{self.folder_path}/{self.txt_name}.txt'
        self.generate_audio_time_total = 0
        self.jobs=[]

    def load_config(self, config_path='config.json'):
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        return config

    def get_folder_path(self):
        return self.folder_path

    def parse_data(self, lines):
        data = []
        for line in lines:
            line = line.strip()
            if line:
                line = line + self.config['stop_suffix'].strip()
                if line == '---':
                    continue
                row = dict()
                row['text'] = line
                row['role'] = 'default'

                pattern = r'%([^%]+)%'
                match = re.search(pattern, line)
                if match:
                    role_name = match.group(1)
                    row['role'] = role_name
                    row['moji'] = ''
                    row['text'] = line.replace(f'%{role_name}%', '')

                    if role_name.count('-') > 0:
                        row['role'] = role_name.split('-')[0]
                        if len(role_name.split('-')) > 1:
                            row['moji'] = role_name.split('-')[1]
                data.append(row)

        return data

    def get_blocks(self, rows):
        blocks = []
        blockContent = ''
        for num, row in enumerate(rows, start=1):
            if row['text'] == '---':
                continue
            if row['text'].startswith("："):
                row['text'] = row['text'][1:]

            if row['role'] != 'default':
                if blockContent:
                    blocks.append({'text': blockContent, 'role': 'default', 'moji': ''})
                blocks.append({'text': row['text'], 'role': row['role'], 'moji': row['moji']})
                blockContent = ''
                continue
            if len(blockContent + row['text']) >= self.config['block_size_limit']:
                blocks.append({'text': blockContent + row['text'], 'role': row['role'], 'moji': ''})
                blockContent = ''
            else:
                blockContent += row['text']

            if num == len(rows):
                moji = ''
                text = ''
                if row['role'] != 'default':
                    moji = row['moji']
                    text = row['text']
                blocks.append({'text': blockContent + text, 'role': row['role'], 'moji': moji})
                blockContent = ''

        return blocks

    def read_file(self, file_path=None):
        if file_path is None:
            file_path = self.file_full_path
        with open(file_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
        return lines

    def save_audio(self, audio_path, name, result):
        dst_file = f'{audio_path}/{name}.wav'
        shutil.copy(result, dst_file)

    def delete_all_files(self, folder_path=None, start_segment=0):
        if folder_path is None:
            folder_path = self.audio_path
        for root, dirs, files in os.walk(folder_path, topdown=False):
            index = 0
            for name in files:
                if index < -start_segment:
                    continue
                index += 1
                file_path = os.path.join(root, name)
                os.remove(file_path)
            for name in dirs:
                dir_path = os.path.join(root, name)
                os.rmdir(dir_path)

        print(f"文件夹 {folder_path} 中的所有文件和子文件夹已被删除")

    def geneAudio(self, block, index, client, use_moji=True):
        if self.beforeGeneAudio(block, index, client) is None:
            return None
        voice, speed = self.beforeGeneAudio(block, index, client)
        seed = 34751218

        start_time = time.time()
        print('speed', speed, block['role'], block['moji'], voice, block['text'])
        reference_audio_path = self.config['reference_audio_path']
        prompt_text = self.config['voices'][voice].split('.')[0]
        if not use_moji or voice in self.config['expect_voices'] or not self.config['moji_enable'] or not block['moji']:
            result = client.predict(
                tts_text=block['text'],
                mode_checkbox_group="预训练音色",
                sft_dropdown=voice,
                prompt_text=prompt_text,
                prompt_wav_upload=handle_file(f'{reference_audio_path}/{self.config["voices"][voice]}'),
                prompt_wav_record=None,
                instruct_text=block['moji'],
                seed=seed,
                stream=False,
                speed=speed,
                api_name="/generate_audio"
            )
        else:
            result = client.predict(
                tts_text=block['text'],
                mode_checkbox_group="自然语言控制",
                sft_dropdown='',
                prompt_text=prompt_text,
                prompt_wav_upload=handle_file(f'{reference_audio_path}/{self.config["voices"][voice]}'),
                prompt_wav_record=None,
                instruct_text=block['moji'],
                seed=seed,
                stream=False,
                speed=speed,
                api_name="/generate_audio"
            )

        generate_audio_time = time.time() - start_time
        self.generate_audio_time_total += generate_audio_time
        print(f"生成音频的时间：{generate_audio_time} 秒")

        dst_file = f'{self.audio_path}/a{index}.wav'
        shutil.copy(result, dst_file)

        return index

    def beforeGeneAudio(self, block, index, client):
        if len(self.config["only_roles"]) > 0 and block['role'] not in self.config["only_roles"]:
            return None
        if block['role'] not in self.config['roles'].keys():
            voice = self.config['default_voice']
        else:
            voice = self.config['roles'][block['role']]
        speed = self.config['speeds']['speaker']

        if block['role'] in self.config['speeds'].keys():
            speed = self.config['speeds'][block['role']]
        return voice,speed


    def geneAudioAsync(self, block, index, client):
        if self.beforeGeneAudio(block, index, client) is None:
            return None
        voice, speed = self.beforeGeneAudio(block, index, client)
        seed = 0
        print('speed', speed, block['moji'], voice)
        reference_audio_path = self.config['reference_audio_path']
        prompt_text = self.config['voices'][voice].split('.')[0]
        if voice in self.config['expect_voices'] or not self.config['moji_enable'] or not block['moji']:
            job = client.submit(
                tts_text=block['text'],
                mode_checkbox_group="预训练音色",
                sft_dropdown=voice,
                prompt_text=prompt_text,
                prompt_wav_upload=handle_file(f'{reference_audio_path}/{self.config["voices"][voice]}'),
                prompt_wav_record=None,
                instruct_text=block['moji'],
                seed=seed,
                stream=False,
                speed=speed,
                api_name="/generate_audio"
            )
        else:
            job = client.submit(
                tts_text=block['text'],
                mode_checkbox_group="自然语言控制",
                sft_dropdown='',
                prompt_text=prompt_text,
                prompt_wav_upload=handle_file(f'{reference_audio_path}/{self.config["voices"][voice]}'),
                prompt_wav_record=None,
                instruct_text=block['moji'],
                seed=seed,
                stream=False,
                speed=speed,
                api_name="/generate_audio"
            )

        self.jobs.append({"index": index, "job": job})
        return len(self.jobs)

    def save_audio_async(self, job, index):
        start_time = time.time()
        print('save_audio_async', job,index)
        result = job.result()
        generate_audio_time = time.time() - start_time
        self.generate_audio_time_total += generate_audio_time
        print(f"生成音频的时间：{generate_audio_time} 秒")

        dst_file = f'{self.audio_path}/a{index}.wav'
        shutil.copy(result, dst_file)

    def generate_audio_run(self, blocks, start_segment, client):
        for i, block in enumerate(blocks, start=0):
            if self.config['test_block_num'] > 0:
                if i > self.config['test_block_num']:
                    break
            if i < start_segment:
                continue
            print(i, len(blocks), len(block['text']), block)
            job_num = self.geneAudioAsync(block, i, client)

            if job_num > self.config['max_job_num']:
                job = self.jobs.pop(0)
                self.save_audio_async(job, i - self.config['max_job_num'])

    def get_role_dict(self, blocks):
        role_dict = {}

        for i, block in enumerate(blocks, start=0):
            if block['role'] not in role_dict.keys():
                role_dict[block['role']] = []
            role_dict[block['role']].append({'index': i, 'block': block})
        return role_dict